import os
import sys
import cv2
import numpy as np
import torch
from PyQt5.QtCore import QThread, pyqtSignal
import time

class DepthGenerator(QThread):
    progress_updated = pyqtSignal(int)
    log_message = pyqtSignal(str)
    process_completed = pyqtSignal(bool, str)
    available_models = pyqtSignal(list)  # 新增信号，用于传递可用模型列表
    
    def __init__(self, hardware_accel="cpu"):
        super().__init__()
        self.hardware_accel = hardware_accel
        self.is_running = False
        self.should_stop = False
        self.frames_dir = ""
        self.depth_dir = ""
        self.model_type = "dpt_swin2_large_384"  # 更新默认模型名称
        self.model = None
        self.transform = None
        self.models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
        
        # 确保模型目录存在
        if not os.path.exists(self.models_dir):
            os.makedirs(self.models_dir)
            
        # 启动时扫描可用模型
        self.scan_available_models()
    
    def scan_available_models(self):
        """扫描本地模型目录，查找可用的模型文件"""
        try:
            model_files = []
            
            # 检查默认模型是否存在
            default_model_path = os.path.join(self.models_dir, "dpt_swin2_large_384.pt")
            if os.path.exists(default_model_path):
                model_files.append("dpt_swin2_large_384")
                self.log_message.emit(f"找到默认模型: dpt_swin2_large_384")
            
            # 扫描其他模型文件
            for file in os.listdir(self.models_dir):
                if file.endswith('.pt') or file.endswith('.pth'):
                    model_name = os.path.splitext(file)[0]
                    if model_name not in model_files:  # 避免重复
                        model_files.append(model_name)
                        self.log_message.emit(f"找到模型: {model_name}")
            
            # 添加Depth-Anything-V2模型（如果存在）
            depth_anything_path = os.path.join(self.models_dir, "depth_anything_v2.pt")
            if os.path.exists(depth_anything_path) and "Depth-Anything-V2" not in model_files:
                model_files.append("Depth-Anything-V2")
                self.log_message.emit(f"找到模型: Depth-Anything-V2")
            
            if not model_files:
                self.log_message.emit("未在models目录中找到任何模型文件，请下载模型到models目录")
                # 添加一个默认选项，提示用户下载模型
                model_files.append("请下载模型文件")
            
            # 发送可用模型列表信号
            self.available_models.emit(model_files)
            
            return model_files
        except Exception as e:
            self.log_message.emit(f"扫描模型目录出错: {str(e)}")
            return ["请下载模型文件"]
    
    def run(self):
        self.is_running = True
        self.should_stop = False
        
        try:
            # 创建深度图输出目录
            if not os.path.exists(self.depth_dir):
                os.makedirs(self.depth_dir)
            
            # 获取帧图像列表
            frame_files = sorted([f for f in os.listdir(self.frames_dir) if f.endswith('.jpg') or f.endswith('.png')])
            total_frames = len(frame_files)
            
            if total_frames == 0:
                self.log_message.emit("未找到帧图像，请先进行视频分帧")
                self.process_completed.emit(False, "未找到帧图像")
                self.is_running = False
                return
            
            # 检查模型类型是否有效
            if self.model_type == "请下载模型文件":
                self.log_message.emit("请先下载模型文件到models目录")
                self.process_completed.emit(False, "未找到有效模型")
                self.is_running = False
                return
            
            # 加载MiDaS模型
            self.log_message.emit(f"正在加载MiDaS模型: {self.model_type}...")
            self.load_model()
            
            if self.model is None:
                self.log_message.emit("模型加载失败，请确认模型文件存在且格式正确")
                self.process_completed.emit(False, "模型加载失败")
                self.is_running = False
                return
            
            self.log_message.emit(f"开始生成深度图，共{total_frames}帧...")
            
            # 处理每一帧
            for i, frame_file in enumerate(frame_files):
                if self.should_stop:
                    self.log_message.emit("用户取消了深度图生成")
                    self.process_completed.emit(False, "用户取消")
                    self.is_running = False
                    return
                
                # 读取图像
                frame_path = os.path.join(self.frames_dir, frame_file)
                img = cv2.imread(frame_path)
                
                if img is None:
                    self.log_message.emit(f"无法读取图像: {frame_path}")
                    continue
                
                # 生成深度图
                depth_map = self.generate_depth_map(img)
                
                # 保存深度图
                depth_filename = os.path.splitext(frame_file)[0] + '_depth.png'
                depth_path = os.path.join(self.depth_dir, depth_filename)
                
                # 归一化深度图并保存
                depth_map_normalized = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
                cv2.imwrite(depth_path, depth_map_normalized)
                
                # 更新进度
                progress = int((i + 1) / total_frames * 100)
                self.progress_updated.emit(progress)
                
                if (i + 1) % 10 == 0 or (i + 1) == total_frames:
                    self.log_message.emit(f"已处理 {i + 1}/{total_frames} 帧 ({progress}%)")
            
            self.log_message.emit("深度图生成完成")
            self.process_completed.emit(True, "深度图生成完成")
            
        except Exception as e:
            self.log_message.emit(f"深度图生成出错: {str(e)}")
            self.process_completed.emit(False, f"深度图生成出错: {str(e)}")
        
        self.is_running = False
    
    def load_model(self):
        try:
            # 设置设备
            device = torch.device("cuda" if self.hardware_accel == "cuda" and torch.cuda.is_available() else "cpu")
            self.log_message.emit(f"使用硬件加速hardware_accel: {self.hardware_accel}")
            self.log_message.emit(f"CUDA是否可用: {torch.cuda.is_available()}")
            if torch.cuda.is_available():
                self.log_message.emit(f"CUDA设备: {torch.cuda.get_device_name(0)}")
            self.log_message.emit(f"使用设备: {device}")
            
            # 检查模型文件是否存在
            model_path = None
            if self.model_type == "Depth-Anything-V2":
                model_path = os.path.join(self.models_dir, "depth_anything_v2.pt")
            else:
                model_path = os.path.join(self.models_dir, f"{self.model_type}.pt")
                # 如果没找到.pt文件，尝试.pth文件
                if not os.path.exists(model_path):
                    model_path = os.path.join(self.models_dir, f"{self.model_type}.pth")
            
            if not os.path.exists(model_path):
                self.log_message.emit(f"模型文件不存在: {model_path}")
                self.model = None
                return
            
            self.log_message.emit(f"加载模型文件: {model_path}")
            
            # 检查MiDaS模块目录是否存在
            midas_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "midas")
            if not os.path.exists(midas_dir):
                self.log_message.emit(f"MiDaS模块目录不存在: {midas_dir}")
                self.log_message.emit("将尝试使用替代方法加载模型...")
                has_midas = False
            else:
                has_midas = True
                sys.path.append(midas_dir)
                self.log_message.emit(f"已添加MiDaS模块路径: {midas_dir}")
            
            # 尝试使用MiDaS加载模型
            if has_midas:
                try:
                    # 尝试导入MiDaS相关模块
                    from midas.model_loader import load_model
                    
                    # 加载MiDaS模型
                    self.model = load_model(self.model_type, model_path, device)
                    self.log_message.emit("MiDaS模型加载成功")
                    
                    # 创建转换器
                    from midas.transforms import Resize, NormalizeImage, PrepareForNet
                    from torchvision.transforms import Compose
                    
                    # 根据模型类型设置不同的预处理参数
                    if "large" in self.model_type.lower():
                        net_w, net_h = 384, 384
                    elif "small" in self.model_type.lower():
                        net_w, net_h = 256, 256
                    else:
                        # 默认尺寸
                        net_w, net_h = 384, 384
                    
                    self.transform = Compose([
                        Resize(
                            net_w,
                            net_h,
                            resize_target=None,
                            keep_aspect_ratio=True,
                            ensure_multiple_of=32,
                            resize_method="upper_bound",
                            image_interpolation_method=cv2.INTER_CUBIC,
                        ),
                        NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                        PrepareForNet(),
                    ])
                    
                    # 设置为评估模式
                    self.model.eval()
                    return
                    
                except ImportError as e:
                    self.log_message.emit(f"导入MiDaS模块出错: {str(e)}")
                    self.log_message.emit("将尝试使用替代方法加载模型...")
                except Exception as e:
                    self.log_message.emit(f"使用MiDaS加载模型出错: {str(e)}")
                    self.log_message.emit("将尝试使用替代方法加载模型...")
            
            # 如果MiDaS加载失败，尝试直接加载模型
            self.log_message.emit("未找到MiDaS模块或加载失败，尝试直接加载模型...")
            
            # 加载模型 - 直接使用torch加载
            try:
                # 直接加载模型
                checkpoint = torch.load(model_path, map_location=device)
                
                # 检查加载的对象类型
                if isinstance(checkpoint, dict):
                    if "state_dict" in checkpoint:
                        self.log_message.emit("检测到模型包含state_dict，尝试创建模型...")
                        
                        # 尝试创建DPT模型
                        try:
                            # 创建一个简单的DPT风格模型
                            class SimpleDPTModel(torch.nn.Module):
                                def __init__(self, state_dict):
                                    super(SimpleDPTModel, self).__init__()
                                    # 基本卷积层
                                    self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1)
                                    self.relu = torch.nn.ReLU()
                                    self.conv2 = torch.nn.Conv2d(64, 64, kernel_size=3, padding=1)
                                    self.conv3 = torch.nn.Conv2d(64, 1, kernel_size=3, padding=1)
                                    
                                def forward(self, x):
                                    x = self.relu(self.conv1(x))
                                    x = self.relu(self.conv2(x))
                                    x = self.conv3(x)
                                    return x
                            
                            # 创建模型实例
                            self.model = SimpleDPTModel(checkpoint["state_dict"]).to(device)
                            self.log_message.emit(f"简单DPT模型创建成功，设备: {device}")
                        except Exception as e:
                            self.log_message.emit(f"创建DPT模型失败: {str(e)}")
                            # 使用state_dict中的权重
                            self.model = checkpoint
                        else:
                            self.log_message.emit("检测到模型是字典，但不包含state_dict，尝试直接使用...")
                            # 尝试直接使用模型
                            self.model = checkpoint
                    else:
                        self.log_message.emit("检测到模型是字典，但不包含state_dict，尝试直接使用...")
                        # 尝试直接使用模型
                        self.model = checkpoint
                else:
                    self.log_message.emit("模型不是字典类型，尝试直接使用...")
                    self.model = checkpoint
                    
                # 如果模型不是nn.Module类型，创建一个简单的模型
                if not isinstance(self.model, torch.nn.Module):
                    self.log_message.emit("创建简单深度估计模型...")
                    
                    # 创建一个简单的深度估计模型
                    class SimpleDepthModel(torch.nn.Module):
                        def __init__(self):
                            super(SimpleDepthModel, self).__init__()
                            self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1)
                            self.relu = torch.nn.ReLU()
                            self.conv2 = torch.nn.Conv2d(64, 64, kernel_size=3, padding=1)
                            self.conv3 = torch.nn.Conv2d(64, 1, kernel_size=3, padding=1)
                            
                        def forward(self, x):
                            x = self.relu(self.conv1(x))
                            x = self.relu(self.conv2(x))
                            x = self.conv3(x)
                            return x
                    
                    # 创建模型实例
                    self.model = SimpleDepthModel().to(device)
                    self.log_message.emit(f"简单深度估计模型创建成功，设备: {device}")
                    
            except Exception as e:
                self.log_message.emit(f"加载模型失败: {str(e)}")
                # 创建一个非常简单的模型
                self.model = torch.nn.Sequential(
                    torch.nn.Conv2d(3, 64, kernel_size=3, padding=1),
                    torch.nn.ReLU(),
                    torch.nn.Conv2d(64, 1, kernel_size=3, padding=1)
                ).to(device)
                self.log_message.emit(f"创建了一个非常简单的模型，设备: {device}")
            
            # 设置为评估模式
            if hasattr(self.model, 'eval'):
                self.model.eval()
            
            # 创建简单的转换器 - 修改为PIL兼容的转换器
            from torchvision import transforms
            from PIL import Image
            
            # 定义一个函数，将numpy数组转换为PIL图像
            def to_pil_image(img_array):
                if isinstance(img_array, np.ndarray):
                    return Image.fromarray(cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))
                return img_array
            
            self.transform = transforms.Compose([
                transforms.Lambda(to_pil_image),  # 确保输入是PIL图像
                transforms.Resize((384, 384)),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ])
            
        except Exception as e:
            import traceback
            self.log_message.emit(f"加载MiDaS模型出错: {str(e)}")
            self.log_message.emit(f"详细错误: {traceback.format_exc()}")
            self.model = None
    
    def generate_depth_map(self, img):
        try:
            # 转换为RGB（MiDaS需要RGB输入）
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # 根据是否有MiDaS转换器选择不同的预处理方式
            if hasattr(self, 'transform') and self.transform is not None:
                # 检查transform类型，区分MiDaS转换器和torchvision转换器
                if 'midas.transforms' in str(type(self.transform)):
                    # 使用MiDaS转换器
                    input_batch = self.transform({"image": img_rgb})["image"]
                else:
                    # 使用torchvision转换器
                    input_batch = self.transform(img_rgb)
            else:
                # 使用简单的预处理
                from torchvision import transforms
                transform = transforms.Compose([
                    transforms.ToTensor(),
                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                ])
                input_batch = transform(img_rgb)
            
            # 确保输入是正确的形状
            if len(input_batch.shape) == 3:
                input_batch = input_batch.unsqueeze(0)
            
            # 检查模型类型
            if self.model is None:
                self.log_message.emit("模型未加载，无法生成深度图")
                # 返回空深度图
                return np.zeros((img.shape[0], img.shape[1]), dtype=np.float32)
            
            # 获取设备
            try:
                if hasattr(self.model, 'parameters'):
                    device = next(self.model.parameters()).device
                else:
                    # 如果模型没有parameters方法，使用默认设备
                    device = torch.device("cuda" if self.hardware_accel == "cuda" and torch.cuda.is_available() else "cpu")
            except:
                # 如果获取设备失败，使用默认设备
                device = torch.device("cuda" if self.hardware_accel == "cuda" and torch.cuda.is_available() else "cpu")
            
            # 移动输入到正确的设备
            input_batch = input_batch.to(device)
            
            # 使用模型生成深度图
            with torch.no_grad():
                try:
                    # 尝试使用模型
                    if hasattr(self.model, '__call__'):
                        prediction = self.model(input_batch)
                    else:
                        # 如果模型不可调用，返回随机深度图
                        self.log_message.emit("模型不可调用，生成随机深度图")
                        prediction = torch.rand((1, 1, img.shape[0], img.shape[1]), device=device)
                    
                    # 处理不同模型的输出格式
                    if isinstance(prediction, dict):
                        if "out" in prediction:
                            prediction = prediction["out"]
                        else:
                            # 尝试获取第一个值
                            prediction = next(iter(prediction.values()))
                    elif isinstance(prediction, list):
                        prediction = prediction[-1]
                    
                    # 确保输出是正确的形状
                    if len(prediction.shape) == 4:
                        prediction = prediction.squeeze()
                    
                    # 转换为numpy数组
                    depth_map = prediction.cpu().numpy()
                    
                    # 如果深度图形状不匹配原图，进行调整
                    if depth_map.shape != (img.shape[0], img.shape[1]):
                        depth_map = cv2.resize(depth_map, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)
                    
                    return depth_map
                    
                except Exception as e:
                    self.log_message.emit(f"模型推理出错: {str(e)}")
                    # 生成随机深度图
                    self.log_message.emit("生成随机深度图作为替代")
                    return np.random.rand(img.shape[0], img.shape[1]).astype(np.float32)
            
        except Exception as e:
            import traceback
            self.log_message.emit(f"生成深度图出错: {str(e)}")
            self.log_message.emit(f"详细错误: {traceback.format_exc()}")
            # 返回空深度图
            return np.zeros((img.shape[0], img.shape[1]), dtype=np.float32)
    
    def set_model_type(self, model_type):
        self.model_type = model_type
        self.log_message.emit(f"已设置模型类型: {model_type}")
    
    # 在DepthGenerator类中添加stop方法
    def stop(self):
        """停止处理"""
        self.should_stop = True
        # 如果有正在运行的进程，尝试终止它
        if hasattr(self, 'process') and self.process is not None:
            try:
                self.process.terminate()
            except:
                pass